#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author: zyx
@Date: 2024/12/08 11:08
@FileName: tools.py
@Description: langchian agent 的工具
"""
from langchain.agents import tool

from langchain_community.utilities import GoogleSerperAPIWrapper
from dotenv import load_dotenv

load_dotenv()


@tool
def search_tool(query: str) -> str:
    """只有在查询实时的信息或者是不知道的事情的时候才会用到这个工具。"""
    serper = GoogleSerperAPIWrapper()
    return serper.run(query)


def create_vector_query_tool(rag_dir) -> str:
    from rag_wrapper import ChromaVectorRAGWrapper

    rag = ChromaVectorRAGWrapper(rag_dir)

    @tool
    def vector_query(query: str) -> str:
        """查询与化学知识库相关的内容时使用这个工具，使用本地向量库对文档进行搜索。"""
        return rag.query(query)

    return vector_query


def create_lightrag_tool(rag_dir) -> str:
    from rag_wrapper import LightRAGWrapper

    rag = LightRAGWrapper(rag_dir)

    @tool
    def lightrag(query: str) -> str:
        """查询与化学知识库相关的内容时使用这个工具，使用lightrag对文档进行搜索。"""
        return rag.query(query)

    return lightrag


def get_rag_tool_by_type(rag_type: str, rag_dir: str):
    match rag_type:
        case "vector_query":
            return create_vector_query_tool(rag_dir)
        case "lightrag":
            return create_lightrag_tool(rag_dir)
        case _:
            raise ValueError(f"不支持的 RAG 类型: {rag_type}")
